import tensorflow.keras.backend as K
import tensorflow as tf

from keras.api._v2.keras.layers import Layer

import nlp_tools

class LabelMaskLoss(Layer):
    def __init__(self,loss_func):
        self.loss_func = loss_func


    def call(self, y_true, y_pred):
        pass

    @tf.autograph.experimental.do_not_convert
    def loss(self,y_true, y_pred):
        #mask = tf.math.logical_not(tf.math.equal(self.input_mask, 0))  # 将y_true 中所有为0的找出来，标记为False
        loss_ = self.loss_func(y_true, y_pred)
        mask = tf.cast(self.input_mask, dtype=loss_.dtype)  # 将前面统计的是否零转换成1，0的矩阵
        loss_ *= mask  # 将正常计算的loss加上mask的权重，就剔除了padding 0的影响
        loss_ = tf.math.divide_no_nan(tf.reduce_sum(loss_), tf.reduce_sum(mask))
        return loss_






